import getopt
import json
import sys
import time
import traceback
from multiprocessing import Pool

import pandas as pd
from es_pandas import es_pandas

from air_web.config.config import config
from air_web.data_platform import sql_engine
from air_web.web_flask.dal.base_dal import EsBaseDal

es_host = config["ES_HOST"]
es = EsBaseDal(es_host)
ep = es_pandas(es_host)


def write_to_es(df, index_name):
    doc_type = index_name + "_doc"
    ep.to_es(
        df,
        index=index_name,
        doc_type=doc_type,
        thread_count=2,
        chunk_size=5000,
        use_index=False,
        request_timeout=60,
    )


def restore_data_structure(df):
    def add_common_columns(row):
        columns = [
            "cons_no",
            "cons_name",
            "data_date",
            "on5",
            "shi",
            "on7",
            "xian",
            "type_code",
            "type_code_sort",
            "real_cons_no",
            "mp_id",
        ]
        return {c: row.eval(c.upper()) for c in columns}

    def add_binary_columns(data):
        binary_columns = [
            "data_time",
            "p_total",
            "p_kt",
            "p_base",
            "p_std_left",
            "p_std_right",
            "is_day_max",
        ]
        return {c: data.get(c.upper()) for c in binary_columns}

    result = []
    for row in df.itertuples():
        tmp_list = []
        json_binary = row.binary
        binary = json.loads(json_binary)

        tmp_acc = add_common_columns(row)

        for tmp_data in binary:
            tmp_list.append(tmp_acc.update(add_binary_columns(tmp_data)))
        result.extend(tmp_list)

    return pd.DataFrame(result)
